
/*
* Copyright (c) 2008 - 2009 NVIDIA Corporation.  All rights reserved.
*
* NVIDIA Corporation and its licensors retain all intellectual property and proprietary
* rights in and to this software, related documentation and any modifications thereto.
* Any use, reproduction, disclosure or distribution of this software and related
* documentation without an express license agreement from NVIDIA Corporation is strictly
* prohibited.
*
* TO THE MAXIMUM EXTENT PERMITTED BY APPLICABLE LAW, THIS SOFTWARE IS PROVIDED *AS IS*
* AND NVIDIA AND ITS SUPPLIERS DISCLAIM ALL WARRANTIES, EITHER EXPRESS OR IMPLIED,
* INCLUDING, BUT NOT LIMITED TO, IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
* PARTICULAR PURPOSE.  IN NO EVENT SHALL NVIDIA OR ITS SUPPLIERS BE LIABLE FOR ANY
* SPECIAL, INCIDENTAL, INDIRECT, OR CONSEQUENTIAL DAMAGES WHATSOEVER (INCLUDING, WITHOUT
* LIMITATION, DAMAGES FOR LOSS OF BUSINESS PROFITS, BUSINESS INTERRUPTION, LOSS OF
* BUSINESS INFORMATION, OR ANY OTHER PECUNIARY LOSS) ARISING OUT OF THE USE OF OR
* INABILITY TO USE THIS SOFTWARE, EVEN IF NVIDIA HAS BEEN ADVISED OF THE POSSIBILITY OF
* SUCH DAMAGES
*/

#include <optix.h>
#include <optixu/optixu_math_namespace.h>
#include "helpers.h"
#include "path_tracer.h"
#include "materials.h"
#include "random.h"
#include "sunsky.h"

using namespace optix;


// Scene wide
rtDeclareVariable(float,         scene_epsilon, , );
rtDeclareVariable(rtObject,      top_object, , );

// For camera
rtDeclareVariable(float3,        eye, , );
rtDeclareVariable(float3,        U, , ); 
rtDeclareVariable(float3,        V, , );
rtDeclareVariable(float3,        W, , );
rtDeclareVariable(float3,        bad_color, , );
rtDeclareVariable(unsigned int,  frame_number, , );
rtDeclareVariable(unsigned int,  sqrt_num_samples, , );

rtBuffer<TriangleLight>     lights;

//ray types
rtDeclareVariable(unsigned int,  pathtrace_ray_type, , );
rtDeclareVariable(unsigned int,  pathtrace_shadow_ray_type, , );
rtDeclareVariable(unsigned int,  rr_begin_depth, , );
rtDeclareVariable(unsigned int,  max_depth, , );
rtDeclareVariable(int, lgt_idx, attribute lgt_idx, ); 
rtDeclareVariable(float,        Y_log_av, , );
rtDeclareVariable(float,        Y_max, , );

//output buffers
rtBuffer<float4, 2>              output_buffer;
rtBuffer<float3, 2>              hdr_buffer;

//rays
rtDeclareVariable(PerRayData_pathtrace, current_prd, rtPayload, );

//optix tracked data
rtDeclareVariable(optix::Ray,	ray,          rtCurrentRay, );
rtDeclareVariable(float,		t_hit,        rtIntersectionDistance, );
rtDeclareVariable(uint2,		launch_index, rtLaunchIndex, );
rtDeclareVariable(uint2,		launch_dim,   rtLaunchDim, );

//geometry attributes
rtDeclareVariable(float3, geometric_normal, attribute geometric_normal, ); 
rtDeclareVariable(float3, shading_normal,   attribute shading_normal, ); 
rtDeclareVariable(float3, texcoord, attribute texcoord, ); 


// For miss program
rtDeclareVariable(float3,       bg_color, , );

//-----------------------------------------------------------------------------
//
//  Camera program -- main ray tracing loop
//
//-----------------------------------------------------------------------------
rtTextureSampler<float4, 2> diffuse_map;
rtTextureSampler<float4, 2> envmap;

__device__ __inline__ optix::float4 rect_envmap_le(optix::float3& dir)
{
  float3 direction = -dir;
  float theta = atan2f( direction.x, direction.y );
  float phi   = M_PIf * 0.5f -  acosf( direction.z );
  float u     = (theta + M_PIf) * (0.5f * M_1_PIf);
  float v     = 0.5f * ( 1.0f + sin(phi) );
  return tex2D(envmap, u, v);
} 

RT_PROGRAM void pathtrace_camera()
{
	size_t2 screen = output_buffer.size();

	float2 inv_screen = 1.0f/make_float2(screen) * 2.f;
	float2 pixel = (make_float2(launch_index)) * inv_screen - 1.f;

	float2 jitter_scale = inv_screen / sqrt_num_samples;
	unsigned int samples_per_pixel = sqrt_num_samples*sqrt_num_samples;
	float3 result = make_float3(0.0f);

	unsigned int seed = tea<16>(screen.x*launch_index.y+launch_index.x, frame_number);

	do {
		unsigned int x = samples_per_pixel%sqrt_num_samples;
		unsigned int y = samples_per_pixel/sqrt_num_samples;
		float2 jitter = make_float2(x-rnd(seed), y-rnd(seed));
		float2 d = pixel + jitter*jitter_scale;
		float3 ray_origin = eye;
		float3 ray_direction = normalize(d.x*U + d.y*V + W);

		PerRayData_pathtrace prd;
		prd.result = make_float3(0.f);
		prd.attenuation = make_float3(1.f);
		prd.countEmitted = true;
		prd.done = false;
		prd.inside = false;
		prd.seed = seed;
		prd.depth = 0;

		for(;;) {
			Ray ray = make_Ray(ray_origin, ray_direction, pathtrace_ray_type, scene_epsilon, RT_DEFAULT_MAX);
			rtTrace(top_object, ray, prd);
			if(prd.done) {
				prd.result += prd.radiance * prd.attenuation;
				break;
			}

			// RR
			if(prd.depth >= rr_begin_depth){
				//break;
				float pcont = fmaxf(prd.attenuation);
				if(rnd(prd.seed) >= pcont || prd.depth > max_depth )
					break;
				prd.attenuation /= pcont;
			}
			prd.depth++;
			prd.result += prd.radiance * prd.attenuation;
			ray_origin = prd.origin;
			ray_direction = prd.direction;
		} // eye ray

		result += prd.result;
		seed = prd.seed;
	} while (--samples_per_pixel);

	float3 pixel_color = result/(sqrt_num_samples*sqrt_num_samples);

	if (frame_number > 1)
	{
		 float a = 1.0f / (float)frame_number;
		float b = ((float)frame_number - 1.0f) * a;
		hdr_buffer[launch_index] = a * pixel_color + b * hdr_buffer[launch_index];
	}
	else
	{
		hdr_buffer[launch_index] = pixel_color;
	}
	output_buffer[launch_index] = make_float4(hdr_buffer[launch_index], 0.0f);
}


rtDeclareVariable(float, time_view_scale, , ) = 1e-6f;

rtDeclareVariable(float, aperture_radius, , );
rtDeclareVariable(float, focal_scale, , );


RT_PROGRAM void pathtrace_tlcamera()
{
	size_t2 screen = output_buffer.size();

	float2 inv_screen = 1.0f/make_float2(screen) * 2.f;
	float2 pixel = (make_float2(launch_index)) * inv_screen - 1.f;

	float2 jitter_scale = inv_screen / sqrt_num_samples;
	unsigned int samples_per_pixel = sqrt_num_samples*sqrt_num_samples;
	float3 result = make_float3(0.0f);

	unsigned int seed = tea<16>(screen.x*launch_index.y+launch_index.x, frame_number);
	do {
		unsigned int x = samples_per_pixel%sqrt_num_samples;
		unsigned int y = samples_per_pixel/sqrt_num_samples;
		float2 jitter = make_float2(x-rnd(seed), y-rnd(seed));
		float2 jt = make_float2(rnd(seed), rnd(seed));

		float2 d = pixel + jitter*jitter_scale;
		float3 ray_origin = eye;
		float3 ray_direction = d.x*U + d.y*V + W;
		float3 ray_target = ray_origin + focal_scale * ray_direction;


		float2 sample = optix::square_to_disk(make_float2(jt.x, jt.y));
		ray_origin = ray_origin + aperture_radius * ( sample.x * normalize( U ) + sample.y * normalize( V ) );
		ray_direction = normalize(ray_target - ray_origin);

		PerRayData_pathtrace prd;
		prd.result = make_float3(0.f);
		prd.attenuation = make_float3(1.f);
		prd.countEmitted = true;
		prd.done = false;
		prd.inside = false;
		prd.seed = seed;
		prd.depth = 0;

		for(;;) {
			Ray ray = make_Ray(ray_origin, ray_direction, pathtrace_ray_type, scene_epsilon, RT_DEFAULT_MAX);
			rtTrace(top_object, ray, prd);
			if(prd.done) {
				prd.result += prd.radiance * prd.attenuation;
				break;
			}

			// RR
			if(prd.depth >= rr_begin_depth){
				//break;
				float pcont = fmaxf(prd.attenuation);
				if(rnd(prd.seed) >= pcont || prd.depth > max_depth )
					break;
				prd.attenuation /= pcont;
			}
			prd.depth++;
			prd.result += prd.radiance * prd.attenuation;
			ray_origin = prd.origin;
			ray_direction = prd.direction;
		} // eye ray

		result += prd.result;
		seed = prd.seed;
	} while (--samples_per_pixel);

	float3 pixel_color = result/(sqrt_num_samples*sqrt_num_samples);

	if (frame_number > 1)
	{
		 float a = 1.0f / (float)frame_number;
		float b = ((float)frame_number - 1.0f) * a;
		hdr_buffer[launch_index] = a * pixel_color + b * hdr_buffer[launch_index];
	}
	else
	{
		hdr_buffer[launch_index] = pixel_color;
	}
	output_buffer[launch_index] = make_float4(
		//hdr_buffer[launch_index], 0.0f);
		tonemap(hdr_buffer[launch_index], Y_log_av, Y_max), 0.0f);
}

rtDeclareVariable(float3,        emission_color, , );
rtDeclareVariable(float3,        diffuse_color, , );
//rtDeclareVariable(PerRayData_pathtrace_bsdf_shadow, bsdf_shadow_prd, rtPayload, );

rtDeclareVariable(float3,       directional_light, , );
rtDeclareVariable(float3,       directional_light_col, , );
rtDeclareVariable(float,      solid_angle, , ) = 6.0e-5f * 100.0f;
rtDeclareVariable(float,		phong_exp, ,);
rtDeclareVariable(float3,        glass_color, , );



RT_PROGRAM void diffuse()
{
	if( length( emission_color ) > 0.0f )
	{
		current_prd.radiance = current_prd.countEmitted? emission_color : make_float3(1.f,1.f,0.f);
		current_prd.done = true;
		return;
	}
	
	float3 world_shading_normal   = normalize( rtTransformNormal( RT_OBJECT_TO_WORLD, shading_normal ) );
	float3 world_geometric_normal = normalize( rtTransformNormal( RT_OBJECT_TO_WORLD, geometric_normal ) );
	float2 uv                     = make_float2(texcoord);
	float3 Kd = make_float3( tex2D(diffuse_map,  uv.x, uv.y) );
	float3 ffnormal = faceforward( world_shading_normal, -ray.direction, world_geometric_normal );

	float3 hitpoint = ray.origin + t_hit * ray.direction;
	current_prd.origin = hitpoint;

	float z1=rnd(current_prd.seed);
	float z2=rnd(current_prd.seed);
	float3 p;
	cosine_sample_hemisphere(z1, z2, p);
	float3 v1, v2;
	createONB(ffnormal, v1, v2);
	//ffnormal = normalize(make_float3( v1.x * mappedNormal.x + v2.x * mappedNormal.y + ffnormal.x * mappedNormal.z, v1.y * mappedNormal.x + v2.y * mappedNormal.y + ffnormal.y * mappedNormal.z, v1.z * mappedNormal.x + v2.z * mappedNormal.y + ffnormal.z * mappedNormal.z));
	current_prd.direction = v1 * p.x + v2 * p.y + ffnormal * p.z;
	current_prd.attenuation = current_prd.attenuation * Kd; // use the diffuse_color as the diffuse response
	current_prd.countEmitted = false;

	unsigned int num_lights = lights.size();
	float3 result = make_float3(0.f);

	//for(int i = 0; i < num_lights; ++i) 
	{
		TriangleLight light = lights[rnd(current_prd.seed)*(num_lights)];
		//TriangleLight light = lights[i];

		float z1 = rnd(current_prd.seed);
		float z2 = rnd(current_prd.seed);
		

		float alpha = 1.0f - sqrt(fold(z1));
		float beta = (1.0f - fold(z2)) * sqrt(fold(z1));
		float gamma = fold(z2) * sqrt(fold(z1));
		float3 light_pos = light.v3 * gamma + light.v1 * alpha + light.v2 * beta;

      float Ldist = length(light_pos - hitpoint);
      float3 L = normalize(light_pos - hitpoint);
      float nDl = dot( ffnormal, L );
      float LnDl = dot( light.normal, L );
      float A = length(cross(light.v2 - light.v3, light.v1 - light.v3));

      // cast shadow ray
      if ( nDl > 0.0f && LnDl > 0.0f ) 
      {
        PerRayData_pathtrace_shadow shadow_prd;
        shadow_prd.inShadow = false;
        Ray shadow_ray = make_Ray( hitpoint, L, pathtrace_shadow_ray_type, scene_epsilon, Ldist );
        rtTrace(top_object, shadow_ray, shadow_prd);

        if (!shadow_prd.inShadow)
        {
          float weight = nDl * LnDl / (M_PIf*Ldist*Ldist) * A;
          result += num_lights * light.emission * weight;
        }
      }
	}
	current_prd.radiance = result;
	//current_prd.result += current_prd.attenuation * result;
	//current_prd.radiance = make_float3(0.0f);
    //current_prd.attenuation = current_prd.attenuation * Kd; 
}


rtDeclareVariable(PerRayData_pathtrace_shadow, current_prd_shadow, rtPayload, );

RT_PROGRAM void shadow()
{
	 if (fmaxf(emission_color) == 0.0f)
  {
    current_prd_shadow.inShadow = true;
	current_prd_shadow.attenuation = make_float3(0.0);
    rtTerminateRay();
  }
}
rtDeclareVariable(float3, shadow_attenuation, , );
rtDeclareVariable(float,        importance_cutoff, , );


RT_PROGRAM void glass_shadow()
{
	float3 world_normal = normalize( rtTransformNormal( RT_OBJECT_TO_WORLD, shading_normal ) );
	float nDi = fabs(dot(world_normal, ray.direction));
	current_prd_shadow.inShadow = true;

	current_prd_shadow.attenuation *= 1-fresnel_schlick(nDi, 5, 1-shadow_attenuation, make_float3(1));
	if(optix::luminance(current_prd_shadow.attenuation) < importance_cutoff) {
		rtTerminateRay();
	}
	else {
		rtIgnoreIntersection();
	}
}



//-----------------------------------------------------------------------------
//
//  Exception program
//
//-----------------------------------------------------------------------------

RT_PROGRAM void exception()
{
	output_buffer[launch_index] = make_float4(bad_color, 0.0f);
}


//-----------------------------------------------------------------------------
//
//  Miss program
//
//-----------------------------------------------------------------------------

RT_PROGRAM void miss()
{
	current_prd.radiance = bg_color;
	current_prd.done = true;
}
